package main

import (
	"fmt"
	"github.com/lib/pq"
	"log"
	"time"
)

var arraySizes []int
var arrayValues [][]string

func init() {

	// 准备数组值 - 按2倍增加 (2,4,8,...,1024)
	arraySizes = []int{2, 4, 8, 16, 32, 64, 128, 256, 512, 1024}
	arrayValues = make([][]string, 10)
	for i := 0; i < 10; i++ {
		arrayValues[i] = make([]string, arraySizes[i])
		for j := 0; j < arraySizes[i]; j++ {
			arrayValues[i][j] = randString(5) // 每个数组元素是5字符的随机字符串
		}
	}
}

func insertTestData(count int) error {
	// 准备固定值
	fixedChar := "fixedValue"
	fixedNum := 12345
	fixedTime := time.Now()

	// 开始事务
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("开始事务失败: %v", err)
	}
	defer tx.Rollback()

	// 准备插入语句
	stmt, err := tx.PrepareContext(ctx, fmt.Sprintf(`
        INSERT INTO %s (
            no_index_char_1, no_index_char_2, no_index_char_3, no_index_char_4, no_index_char_5,
            no_index_char_6, no_index_char_7, no_index_char_8, no_index_char_9, no_index_char_10,
            no_index_char_11, no_index_char_12, no_index_char_13, no_index_char_14, no_index_char_15,
            no_index_char_16, no_index_char_17, no_index_char_18, no_index_char_19, no_index_char_20,
            no_index_char_21, no_index_char_22, no_index_char_23, no_index_char_24, no_index_char_25,
            no_index_char_26, no_index_char_27, no_index_char_28, no_index_char_29, no_index_char_30,
            
            index_char_1, index_char_2, index_char_3, index_char_4, index_char_5,
            index_char_6, index_char_7, index_char_8, index_char_9, index_char_10,
            
            index_num_1, index_num_2, index_num_3, index_num_4, index_num_5,
            index_num_6, index_num_7, index_num_8, index_num_9, index_num_10,
            
            index_time_1, index_time_2, index_time_3, index_time_4, index_time_5,
            index_time_6, index_time_7, index_time_8, index_time_9, index_time_10,
            
            index_array_1, index_array_2, index_array_3, index_array_4, index_array_5,
            index_array_6, index_array_7, index_array_8, index_array_9, index_array_10
        ) VALUES (
            $1, $2, $3, $4, $5, $6, $7, $8, $9, $10,
            $11, $12, $13, $14, $15, $16, $17, $18, $19, $20,
            $21, $22, $23, $24, $25, $26, $27, $28, $29, $30,
            
            $31, $32, $33, $34, $35, $36, $37, $38, $39, $40,
            
            $41, $42, $43, $44, $45, $46, $47, $48, $49, $50,
            
            $51, $52, $53, $54, $55, $56, $57, $58, $59, $60,
            
            $61, $62, $63, $64, $65, $66, $67, $68, $69, $70
        )
    `, tableName))
	if err != nil {
		return fmt.Errorf("准备插入语句失败: %v", err)
	}
	defer stmt.Close()

	// 批量插入数据
	for i := 0; i < count; i++ {
		args := make([]interface{}, 0, 70)

		// 1-30: 无索引字符字段 (30个随机字符串)
		for j := 1; j <= 30; j++ {
			args = append(args, randString(30))
		}

		// 31-40: 有索引字符字段 (10个固定字符串)
		for j := 1; j <= 10; j++ {
			args = append(args, fixedChar)
		}

		// 41-50: 有索引数字字段 (10个固定数字)
		for j := 1; j <= 10; j++ {
			args = append(args, fixedNum)
		}

		// 51-60: 有索引时间字段 (10个固定时间)
		for j := 1; j <= 10; j++ {
			args = append(args, fixedTime)
		}

		// 61-70: 有索引数组字段 (10个不同大小的数组)
		for j := 0; j < 10; j++ {
			args = append(args, pq.Array(arrayValues[j]))
		}

		_, err := stmt.ExecContext(ctx, args...)
		if err != nil {
			return fmt.Errorf("插入数据失败: %v", err)
		}

		if i%1000 == 0 {
			log.Printf("已插入 %d 条记录", i)
		}
	}

	// 提交事务
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("提交事务失败: %v", err)
	}

	log.Printf("成功插入 %d 条测试数据", count)
	return nil
}
